# -*- coding: utf-8 -*-
from datetime import datetime
from collections import Counter
from general.generaltools import Debugger
from general.stringoperations import ngram
from general.mathoperations import normalize_dict
from objects.entity import CandidateEntity

NGRAM_ORDER = 5
KEYPHRASENESS = 0.01

class Term(Debugger):
    def __init__(self, name="", name_nospace="", \
            docid="", beg=-1, end=-1, line_num=-1, line_beg=-1, \
            segmentalized_line_beg=-1, keyphraseness=0, keyword_id='', \
            entity_db={}, keyword_db={}, isDebug=False, mode='std'):
        Debugger.__init__(self, isDebug, mode)
        self.isDebug = isDebug
        self.name = name
        self.name_nospace = name_nospace
        self.docid = docid
        self.beg = beg
        self.end = end
        self.line_num = line_num
        self.line_beg = line_beg
        self.segmentalized_line_beg = segmentalized_line_beg
        self.keyword_id = keyword_id
        self.keyphraseness = keyphraseness
        self.candidate_entities = {}
        self.cand_ent_id_list = []
        self.entity_db = {}
        self.keyword_db = {}
        self.query_db = {}
        self.context = {}

    def __str__(self):
        return "term: {0:<40}\n\tkeyphraseness: {1:<12F}\n\tcandidate_entity:"\
                .format(self.name, self.keyphraseness)

    def init_cand_ent(self, mysql, **argv):
        cand_ent_id_list = list(map(int, argv['entity_list']))
        distr_list = list(map(float, argv['distr_list']))
        self.cand_ent_id_list = cand_ent_id_list
        for i in range(len(cand_ent_id_list)):
            entity_info = {}
            entity_info['entity_id'] = cand_ent_id_list[i]
            entity_info['entity_title'] = self.entity_db[entity_info['entity_id']] \
                    if entity_info['entity_id'] in self.entity_db else "Not IN"
            entity_info['prior'] = distr_list[i]
            entity_info['edit_distance'] = 0.0
            entity_info['issub'] = 0
            cand_ent = CandidateEntity(**entity_info)
            self.candidate_entities[entity_info['entity_id']] = cand_ent
        self.update_cand_info(mysql)
        self.sort_cand_ent_by_prior()

    def update_cand_info(self, mysql):
        for entity_id, inlink_list, infreq_list, outlink_list, outfreq_list, \
                category_list in mysql.get_entity_info_from_mem(*self.cand_ent_id_list):
            current_ent = self.candidate_entities[entity_id]
            current_ent.inlink_list = inlink_list
            #current_ent.second_inlink_list = mysql.get_second_order_inlink(*inlink_list)
            current_ent.infreq_dict = dict(zip(inlink_list, infreq_list))
            current_ent.outlink_list = outlink_list
            current_ent.outfreq_dict = dict(zip(outlink_list, outfreq_list))
            current_ent.category_list = category_list

    def cal_local_features(self, mysql=""):
        if not self.cand_ent_id_list:
            return
        self.tl_1 = datetime.now()
        for entity_id in self.cand_ent_id_list:
            self.candidate_entities[entity_id].calculate_local(self)
        self.tl_2 = datetime.now()
        inlink_v, outlink_v = self.cal_local_concept_vector(mysql)
        self.cal_local_concept_features(inlink_v, outlink_v)
        self.tl_3 = datetime.now()
        category_v = self.cal_local_category_vector(mysql)
        self.cal_local_category_features(category_v)
        self.tl_4 = datetime.now()

        #self.cal_bm25(mysql)

    def cal_bm25(self, mysql):
        for entity_id, tf_dict, fl in mysql.get_tf_dict(*self.cand_ent_id_list):
            self.candidate_entities[entity_id].tf_dict = tf_dict
            self.candidate_entities[entity_id].fl = fl
        for entity_id in self.cand_ent_id_list:
            self.candidate_entities[entity_id].calculate_bm25(self.context)

    def cal_local_concept_vector(self, mysql=""):
        inlink_set = set()
        outlink_set = set()
        inlink_v = {}
        outlink_v = {}
        inlink_dict = {}
        outlink_dict = {}
        inlink_dict_back = {}
        outlink_dict_back = {}

        for entity_id in self.cand_ent_id_list:
            current_ent = self.candidate_entities[entity_id]
            #collect all in&out links from candidate entities
            inlink_set.update(current_ent.inlink_list)
            outlink_set.update(current_ent.outlink_list)

        for entity_id, inlink_list, infreq_list, outlink_list, outfreq_list \
                in mysql.get_entity_link_info_from_mem(*inlink_set):
            inlink_dict[entity_id] = len(outlink_list)
        for entity_id, inlink_list, infreq_list, outlink_list, outfreq_list \
                in mysql.get_entity_link_info_from_mem(*outlink_set):
            outlink_dict[entity_id] = len(inlink_list)

        for entity_id in self.cand_ent_id_list:
            current_ent = self.candidate_entities[entity_id]
            if current_ent.entity_title == "Not IN": continue
            for inlink in current_ent.inlink_list:
                if inlink in inlink_v:
                    inlink_v[inlink] += current_ent.infreq_dict[inlink] * current_ent.weight
                    inlink_dict_back[inlink] += 1
                else:
                    inlink_v[inlink] = current_ent.infreq_dict[inlink] * current_ent.weight
                    inlink_dict_back[inlink] = 1
            for outlink in current_ent.outlink_list:
                if outlink in outlink_v:
                    outlink_v[outlink] += current_ent.outfreq_dict[outlink] * current_ent.weight
                    outlink_dict_back[outlink] += 1
                else:
                    outlink_v[outlink] = current_ent.outfreq_dict[outlink] * current_ent.weight
                    outlink_dict_back[outlink] = 1

        for inlink in inlink_v:
            #if inlink not in inlink_dict: continue
            if inlink_dict[inlink] == 0:
                print("Error---Zero inlink:", inlink)
                for entity_id in self.cand_ent_id_list:
                    current_ent = self.candidate_entities[entity_id]
                    for inn in current_ent.inlink_list:
                        if inn == inlink:
                            print(current_ent)
                inlink_v[inlink] = inlink_v[inlink] / inlink_dict_back[inlink]
            else:
                inlink_v[inlink] = inlink_v[inlink] / inlink_dict[inlink]
        for outlink in outlink_v:
            #if outlink not in outlink_dict: continue
            if outlink_dict[outlink] == 0:
                print("Error---Zero outlink:", outlink)
                for entity_id in self.cand_ent_id_list:
                    current_ent = self.candidate_entities[entity_id]
                    for out in current_ent.outlink_list:
                        if out == outlink:
                            print(current_ent)
                outlink_v[outlink] = outlink_v[outlink] / outlink_dict_back[outlink]
            else:
                outlink_v[outlink] = outlink_v[outlink] / outlink_dict[outlink]

        normalize_dict(inlink_v)
        normalize_dict(outlink_v)

        return inlink_v, outlink_v

    def cal_local_concept_features(self, inlink_v, outlink_v):
        for entity_id in self.cand_ent_id_list:
            current_ent = self.candidate_entities[entity_id]
            if current_ent.entity_title == "Not IN": continue
            current_ent.calculate_local_concept_weight(inlink_v, outlink_v)

    def cal_local_category_vector(self, mysql):
        category_id_set = set()
        category_dict = {}
        category_v = {}
        category_dict_back = {}

        for entity_id, category_list\
                in mysql.get_entity_category_info_from_mem(*self.cand_ent_id_list):
            #collect all in&out links from candidate entities
            category_id_set.update(category_list)
        for category_id, category_freq\
                in mysql.get_category_info_from_mem(*category_id_set):
            category_dict[category_id] = category_freq
        for entity_id in self.cand_ent_id_list:
            current_ent = self.candidate_entities[entity_id]
            if current_ent.entity_title == "Not IN": continue
            for category_id in current_ent.category_list:
                if category_id in category_v:
                    category_v[category_id] += current_ent.weight
                    category_dict_back[category_id] += 1
                else:
                    category_v[category_id] = current_ent.weight
                    category_dict_back[category_id] = 1

        for category_id in category_v:
            if category_dict[category_id] == 0:
                print("Error---Zero category:", category_id)
                for entity_id in self.cand_ent_id_list:
                    current_ent = self.candidate_entities[entity_id]
                    for cate_id in current_ent.category_list:
                        if cate_id == category_id:
                            print(current_ent)
                category_v[category_id] = category_v[category_id] / category_dict_back[category_id]
            else:
                category_v[category_id] = category_v[category_id] / category_dict[category_id]

        normalize_dict(category_v)

        return category_v

    def cal_local_category_features(self, category_v):
        for entity_id in self.cand_ent_id_list:
            current_ent = self.candidate_entities[entity_id]
            if current_ent.entity_title == "Not IN": continue
            current_ent.calculate_local_category_weight(category_v)

    def sort_cand_ent_by_prior(self):
        self.cand_ent_id_list = sorted(self.cand_ent_id_list, \
                key=lambda entity_id: -self.candidate_entities[entity_id].prior)

    def print_result(self):
        if self.isDebug:
            if self.keyword_id:
                print("term: {0:<40}\n\tkeyphraseness: {1:<12F}\n\tcandidate_entity:"\
                        .format(self.name, self.keyphraseness))
                for entity_id in self.cand_ent_id_list:
                    print(self.candidate_entities[entity_id])
                    print()
            else:
                print("term: {0}\n\tkeyphraseness: {1}".format(self.name_nospace, 'NULL'))
            print()
            print()

class Query(Term):
    def __init__(self, queryid="", name="", name_nospace="", name_with_newline="", \
            docid="", beg=-1, end=-1, line_num=-1, line_beg=-1, \
            segmentalized_line_beg=-1, keyphraseness=0, keyword_id='', \
            context_terms=[], isDebug=False, mode='std'):
        Term.__init__(self, name=name, name_nospace=name_nospace, docid=docid, \
                beg=beg, end=end, line_num=line_num, line_beg=line_beg, \
                segmentalized_line_beg=segmentalized_line_beg, \
                keyphraseness=keyphraseness, keyword_id=keyword_id, \
                isDebug=isDebug, mode=mode)
        self.isDebug = isDebug
        self.queryid = queryid
        self.name_with_newline= name_with_newline
        self.context_terms = context_terms.copy()

    def __str__(self):
        return "Query(%s)\n\tname:\t%s\n\tname_nospace:\t%s\n" % \
                (self.queryid, self.name, self.name_nospace) + \
                "\tname_with_newline:\t%s\n\tdocid:\t%s\n\tbeg:\t%d\n\tend:\t%d\n" % \
                (self.name_with_newline, self.docid, self.beg, self.end)

    def update_query_location(self, doc):
        doc.update_query_location(self)
        counter = Counter()
        for line_num in doc.get_query_context(self):
            counter.update(doc.segmentalized_sentences[line_num])
        self.context = dict(counter)
        if self.isDebug:
            doc.print_query_context(self)
            doc.dPrint(doc.bugInfo)

    def identify_terms(self, doc="", mysql=""):
        ngrams = []
        id_list = []
        result = {}
        terms_id_list = []
        self.term_set = set()
        #Get Context of Query
        for line_num in doc.get_query_context(self):
            ngrams += ngram(doc.segmentalized_sentences[line_num], NGRAM_ORDER)
        ngrams.append(self.name_nospace)
        #Get keyword_id of Context words
        for ngram_word in set(ngrams):
            if ngram_word in self.keyword_db:
                id_list.append(str(self.keyword_db[ngram_word]))
        #Get keyphraseness, candidate_entity_id_list, distr_list for each term
        for keyword_id, keyword, keyphraseness in mysql.get_keyphraseness(*id_list):
            if keyphraseness >= KEYPHRASENESS or keyword == self.name_nospace:
                result[keyword_id] = {'keyword_id' : keyword_id, 'keyword' : keyword, \
                        'keyphraseness' : keyphraseness, 'entity_list' : [], 'distr_list' : []}
                terms_id_list.append(str(keyword_id))
        for keyword_id, entity_list, distr_list in mysql.get_kpp_from_mem(*terms_id_list):
            result[keyword_id].update({'entity_list' : entity_list, 'distr_list' : distr_list})
        #Update Query keyword info and candidate_entities info
        if self.name_nospace in self.keyword_db:
            self.keyword_id = self.keyword_db[self.name_nospace]
            self.keyphraseness = result[self.keyword_id]['keyphraseness']
            self.init_cand_ent(mysql, **result[self.keyword_id]) #update candidate_entities
            result.pop(self.keyword_id)
        #Init terms of Query
        for keyword_id in result:
            keyword = result[keyword_id]['keyword']
            self.term_set.add(keyword)
            keyphraseness = result[keyword_id]['keyphraseness']
            if not result[keyword_id]['entity_list']: continue
            term = Term(name=keyword, name_nospace=keyword, keyword_id=keyword_id, docid=self.docid, \
                    keyphraseness=keyphraseness, entity_db=self.entity_db, \
                    keyword_db=self.keyword_db, isDebug=self.isDebug)
            term.context = self.context     #Need further analysis
            term.init_cand_ent(mysql, **result[keyword_id])
            self.context_terms.append(term)
        sorted(self.context_terms, key=lambda term: -term.keyphraseness)

    def get_global_set(self):
        inlink_set = set()
        outlink_set = set()
        category_set = set()
        for term in self.context_terms:
            for entity_id in term.cand_ent_id_list:
                current_ent = term.candidate_entities[entity_id]
                inlink_set.update(current_ent.inlink_list)
                outlink_set.update(current_ent.outlink_list)
                #category_set.update(current_ent.category_list)
        temp_in_set = set()
        temp_out_set = set()
        #temp_cate_set = set()
        for entity_id in self.cand_ent_id_list:
            current_ent = self.candidate_entities[entity_id]
            temp_in_set.update(current_ent.inlink_list)
            temp_out_set.update(current_ent.outlink_list)
            #temp_cate_set.update(current_ent.category_list)
        inlink_set.intersection_update(temp_in_set)
        outlink_set.intersection_update(temp_out_set)
        #category_set.intersection_update(temp_cate_set)
        return inlink_set, outlink_set, category_set

    def cal_global_features(self, mysql=""):
        if not self.cand_ent_id_list:
            return
        self.t1 = datetime.now()

        inlink_set, outlink_set, category_id_set = self.get_global_set()
        self.t2 = datetime.now()
        inlink_v, outlink_v = self.cal_global_concept_vector(inlink_set, \
                outlink_set, mysql)
        self.t3 = datetime.now()
        self.cal_global_concept_features(inlink_v, outlink_v)
        self.t4 = datetime.now()

        #self.cal_global_google_distance(mysql)
        self.t5 = datetime.now()
        #category_v = self.cal_global_category_vector(category_id_set, mysql)
        #self.cal_global_category_features(category_v)

    def cal_global_concept_vector(self, inlink_set, outlink_set, mysql=""):
        inlink_v = {}
        outlink_v = {}
        inlink_dict = {}
        outlink_dict = {}
        inlink_dict_back = {}
        outlink_dict_back = {}

        if inlink_set:
            for entity_id, inlink_list, infreq_list, outlink_list, outfreq_list \
                    in mysql.get_entity_link_info_from_mem(*inlink_set):
                inlink_dict[entity_id] = len(outlink_list)
        if outlink_set:
            for entity_id, inlink_list, infreq_list, outlink_list, outfreq_list \
                    in mysql.get_entity_link_info_from_mem(*outlink_set):
                outlink_dict[entity_id] = len(inlink_list)

        for entity_id in self.cand_ent_id_list:
            current_ent = self.candidate_entities[entity_id]
            if current_ent.entity_title == "Not IN": continue
            if inlink_set:
                for inlink in current_ent.inlink_list:
                    if inlink in inlink_v:
                        inlink_v[inlink] += current_ent.infreq_dict[inlink] * current_ent.weight
                        inlink_dict_back[inlink] += 1
                    else:
                        inlink_v[inlink] = current_ent.infreq_dict[inlink] * current_ent.weight
                        inlink_dict_back[inlink] = 1
            if outlink_set:
                for outlink in current_ent.outlink_list:
                    if outlink in outlink_v:
                        outlink_v[outlink] += current_ent.outfreq_dict[outlink] * current_ent.weight
                        outlink_dict_back[outlink] += 1
                    else:
                        outlink_v[outlink] = current_ent.outfreq_dict[outlink] * current_ent.weight
                        outlink_dict_back[outlink] = 1

        if inlink_set:
            for inlink in inlink_dict:
                if inlink_dict[inlink] == 0:
                    print("Error---Zero inlink:", inlink)
                    for entity_id in self.cand_ent_id_list:
                        current_ent = self.candidate_entities[entity_id]
                        for inn in current_ent.inlink_list:
                            if inn == inlink:
                                print(current_ent)
                    inlink_dict[inlink] = inlink_v[inlink] / inlink_dict_back[inlink]
                else:
                    inlink_dict[inlink] = inlink_v[inlink] / inlink_dict[inlink]
            normalize_dict(inlink_dict)
        if outlink_set:
            for outlink in outlink_dict:
                if outlink_dict[outlink] == 0:
                    print("Error---Zero outlink:", outlink)
                    for entity_id in self.cand_ent_id_list:
                        current_ent = self.candidate_entities[entity_id]
                        for out in current_ent.outlink_list:
                            if out == outlink:
                                print(current_ent)
                    outlink_dict[outlink] = outlink_v[outlink] / outlink_dict_back[outlink]
                else:
                    outlink_dict[outlink] = outlink_v[outlink] / outlink_dict[outlink]
            normalize_dict(outlink_dict)

        return inlink_dict, outlink_dict

    def cal_global_concept_features(self, inlink_v, outlink_v):
        for entity_id in self.cand_ent_id_list:
            current_ent = self.candidate_entities[entity_id]
            if current_ent.entity_title == "Not IN": continue
            current_ent.calculate_global_concept_weight(inlink_v, outlink_v)

    def cal_global_category_vector(self, category_id_set, mysql):
        if not category_id_set:
            return {}
        category_dict = {}
        category_v = {}
        category_dict_back = {}

        for category_id, category_freq\
                in mysql.get_category_info_from_mem(*category_id_set):
            category_dict[category_id] = category_freq
        for entity_id in self.cand_ent_id_list:
            current_ent = self.candidate_entities[entity_id]
            if current_ent.entity_title == "Not IN": continue
            for category_id in current_ent.category_list:
                if category_id in category_v:
                    category_v[category_id] += current_ent.weight
                    category_dict_back[category_id] += 1
                else:
                    category_v[category_id] = current_ent.weight
                    category_dict_back[category_id] = 1

        for category_id in category_dict:
            if category_dict[category_id] == 0:
                print("Error---Zero category:", category_id)
                for entity_id in self.cand_ent_id_list:
                    current_ent = self.candidate_entities[entity_id]
                    for cate_id in current_ent.category_list:
                        if cate_id == category_id:
                            print(current_ent)
                category_dict[category_id] = category_v[category_id] / category_dict_back[category_id]
            else:
                category_dict[category_id] = category_v[category_id] / category_dict[category_id]

        normalize_dict(category_dict)

        return category_dict

    def cal_global_category_features(self, category_v):
        for entity_id in self.cand_ent_id_list:
            current_ent = self.candidate_entities[entity_id]
            if current_ent.entity_title == "Not IN": continue
            current_ent.calculate_global_category_weight(category_v)

    def cal_global_google_distance(self, mysql):
        t = -(datetime.now() - datetime.now())
        term_name = ''
        for term in self.context_terms:
            t1 = datetime.now()
            for term_entity_id in term.cand_ent_id_list:
                term_ent = term.candidate_entities[term_entity_id]
                for entity_id in self.cand_ent_id_list:
                    current_ent = self.candidate_entities[entity_id]
                    ret = current_ent.calculate_first_order_google_distance(term_ent)
                    if ret == -1:
                        if current_ent.second_inlink_list == [0]:
                            current_ent.second_inlink_list = mysql.get_second_order_inlink(*current_ent.inlink_list)
                        if term_ent.second_inlink_list == [0]:
                            term_ent.second_inlink_list = mysql.get_second_order_inlink(*term_ent.inlink_list)
                        current_ent.calculate_second_order_google_distance(term_ent)
            t2 = datetime.now()
            if t2 - t1 > t:
                t = t2 - t1
                term_name = term.name_nospace
        self.t = t
        self.term_name = term_name
        for entity_id in self.cand_ent_id_list:
            self.candidate_entities[entity_id].update_google_distance()

    def print_result(self):
        query_info = self.query_db[self.queryid]
        print(list(self.term_set))
        if self.cand_ent_id_list:
            print('cal_local_str\t:', self.tl_2 - self.tl_1)
            print('cal_local_con\t:', self.tl_3 - self.tl_2)
            print('cal_local_cat\t:', self.tl_4 - self.tl_3)
            print('get_global_set\t:', self.t2 - self.t1)
            print('cal_global_concept_vec\t:', self.t3 - self.t2)
            print('cal_global_concept_fea\t:', self.t4 - self.t3)
            print('cal_google_distance\t:', self.t5 - self.t4)
            #print('\tmax_term\t:', self.term_name)
            #print('\tmax_time\t:', self.t)
        if self.isDebug:
            if self.keyword_id:
                r = "{0:<20}: {1:<40}\n{2:<20}: {3:<}\n{4:<20}: {5:<12F}\n".\
                        format('query', self.name, 'query id', self.queryid, 'keyphraseness', self.keyphraseness) + \
                        "{0:<20}: {1:<}\n{2:<20}: {3:<}\n{4:<20}: {5:<}\n\ncandidate entitites:".\
                        format('entity id', query_info['entity_id'], 'entity_title', query_info['en_title'], 'is in', str(query_info['is_in']))
                print(r)
                for entity_id in self.cand_ent_id_list:
                    print(self.candidate_entities[entity_id])
                    print()
            else:
                r = "{0:<20}: {1:<40}\n{2:<20}: {3:<}\n{4:<20}: {5:<}\n".\
                        format('query', self.name, 'query id', self.queryid, 'keyphraseness', 'NULL') + \
                        "{0:<20}: {1:<}\n{2:<20}: {3:<}\n{4:<20}: {5:<}\n\ncandidate entitites:".\
                        format('entity id', query_info['entity_id'], 'entity_title', query_info['en_title'], 'is in', str(query_info['is_in']))
                print(r)
            print()
            print()
